import os
from datetime import datetime, timedelta

import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from cartopy.feature import NaturalEarthFeature
from ecmwf.opendata import Client

parameters = ['tp']

os.makedirs("./data", exist_ok=True)
filename = './data/aifs_ens_cf_medium_rain_acc.grib'

client = Client("ecmwf", beta=False, model="aifs-ens")

now_utc = datetime.utcnow()

available_time = now_utc - timedelta(hours=5)
available_hour = available_time.hour

# 根据可用时间确定最新的预报时次
if available_hour >= 18:
    time = 18
elif available_hour >= 12:
    time = 12
elif available_hour >= 6:
    time = 6
elif available_hour >= 0:
    time = 0

print(f'https://data.ecmwf.int/forecasts/{available_time.strftime("%Y%m%d")}/{time:02d}z')

try:
    client.retrieve(
        date=available_time.strftime("%Y-%m-%d"),
        time=time,
        step=6,
        stream="enfo",
        type="cf",
        levtype="sfc",
        param=parameters,
        target=filename
    )

    # 使用 xarray 读取 GRIB 文件
    data = xr.open_dataset(filename, engine='cfgrib')

    # 获取降水数据
    tp = data['tp']  # Total precipitation

    # 创建图形
    fig = plt.figure(figsize=(14, 10))
    ax = plt.axes(projection=ccrs.PlateCarree())

    # 设置地图范围
    ax.set_extent([70, 140, 15, 55], crs=ccrs.PlateCarree())
    # ax.set_extent([114, 120, 29, 35], crs=ccrs.PlateCarree())

    # 添加地图特征
    ax.add_feature(cfeature.COASTLINE, color='#2F4F4F', linewidth=0.8)
    ax.add_feature(cfeature.BORDERS, color='#2F4F4F', linewidth=0.5)
    ax.add_feature(cfeature.LAND, color='#F5F5DC', alpha=0.8)
    ax.add_feature(cfeature.OCEAN, color='white', alpha=1.0)

    # 中国省级行政区
    provinces = NaturalEarthFeature(
        category='cultural',
        name='admin_1_states_provinces_lines',
        scale='50m',
        facecolor='none',
        edgecolor='gray',
        linewidth=0.8,
        alpha=0.9
    )
    ax.add_feature(provinces)

    # 定义降水等级和对应颜色（与原metview配色相近）
    levels = [0.1, 1, 2, 5, 10, 15, 20, 30, 40, 50, 100, 300, 1000]
    colors = ['#F0DDB7', '#B5FA6B', '#78F42D', '#B5F0FA', '#78BAFF',
              '#3D97F5', '#1F6BED', '#FFEA78', '#FF9F00', '#FF0000',
              '#A32121', '#FF00FF', '#808080']

    # 创建降水等值线填充图
    cs = ax.contourf(tp.longitude, tp.latitude, tp.values,
                     levels=levels, colors=colors, extend='max',
                     transform=ccrs.PlateCarree())

    # 添加颜色条
    cbar = plt.colorbar(cs, ax=ax, orientation='horizontal',
                        pad=0.08, shrink=0.8, aspect=40)
    cbar.set_label('Total accumulated precipitation (mm)',
                   fontsize=12, color='#2F4F4F')
    cbar.ax.tick_params(labelsize=10, colors='#2F4F4F')

    # 添加网格线
    gl = ax.gridlines(draw_labels=True, dms=True, x_inline=False, y_inline=False,
                      alpha=0.6, color='#D2B48C', linewidth=0.5)
    gl.top_labels = False
    gl.right_labels = False

    # 获取时间信息
    valid_time = data.valid_time.values
    valid_time_beijing = valid_time + np.timedelta64(8, 'h')
    valid_time_str = str(valid_time_beijing).replace('T', ' ').split('.')[0] + ' CST'

    start_time = data.time.values
    start_time_beijing = start_time + np.timedelta64(8, 'h')
    start_time_str = str(start_time_beijing).replace('T', ' ').split('.')[0] + ' CST'

    step = data.step.values
    step_hours = step / np.timedelta64(1, 'h')
    step_str = f"{int(step_hours)}h"

    # 添加标题
    title_text = f'AIFS ENS Control: Total Accumulated Precipitation\n'
    title_text += f'START TIME: {start_time_str}, '
    title_text += f'VALID TIME: {valid_time_str}, STEP: {step_str}'

    plt.suptitle(title_text, fontsize=14, y=0.98, color='#2F4F4F')

    # 添加版权信息
    copyright_text = '© European Centre for Medium-Range Weather Forecasts (ECMWF)'

    plt.figtext(0.5, 0.02, copyright_text, ha='center', fontsize=9,
                color='#2F4F4F', style='italic')

    plt.tight_layout()

    # 保存图片
    # output_filename = 'aifs_precipitation_map.png'
    # plt.savefig(output_filename, dpi=300, bbox_inches='tight',
    #             edgecolor='none')

    plt.show()

except Exception as e:
    print(f"错误: {e}")
